import numpy as np
import matplotlib.pyplot as plt

moive = ['Love Movie', 'Action Movie']

T = np.array([[3, 104, 0],
     [2, 100, 0],
     [1, 81, 0,],
     [101, 10, 1],
     [99, 5, 1],
     [98, 2, 1]])
m = len(T)

x = [18, 90]
k = 3

cmap = plt.cm.get_cmap('rainbow', 2)
idx_neg = T[:, 2] == 0
idx_pos = np.invert(idx_neg)
m_neg = np.sum(idx_neg)
m_pos = np.sum(idx_pos)
c00 = cmap(0)
c0 = np.tile(cmap(0), m_neg)
plt.scatter(T[idx_neg, 0], T[idx_neg, 1], c=np.atleast_2d(cmap(0)), label='negative')
plt.scatter(T[idx_pos, 0], T[idx_pos, 1], c=np.atleast_2d(cmap(1)), label='positive')
plt.scatter(x[0], x[1], c='b')
plt.legend()

dis = np.zeros([m, 2])
for i, t in enumerate(T):
    # 欧氏距离
    d = np.sqrt((t[0] - x[0])**2 + (t[1] - x[1])**2)
    dis[i, 0] = d
    dis[i, 1] = t[2]

arg = np.argsort(dis[:, 0])
dis = dis[arg]
print(dis)
dis = dis[:k]

plt.show()
